# 3d_trajectory_y_axis_spin_comparison.py
# Compares trajectories under Y-Axis spin Magnus Force and Crosswind.

import numpy as np
from math import sin, cos, radians, sqrt, pi
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# --- I. Physical Constants and Projectile Parameters ---
RHO0 = 1.225         # Sea-level Air Density (kg/m^3)
T0 = 288.15          # Sea-level Temperature (K)
L = 0.0065           # Temperature lapse rate (K/m)
G = 9.80665          # Gravity (m/s^2)
R = 287.05           # Gas constant (J/(kg*K))
P_TERM = (G / (R * L)) - 1
C_SOUND = 338.0      # Speed of Sound (m/s)

# Standard Projectile Parameters
M = 5.0              # Mass (kg)
A = 0.0095           # Cross-sectional Area (m^2)
D = 0.11             # Caliber / Diameter (m)
MAGNUS_CL = 0.05     # Magnus Lift Coefficient (simplified constant)

# Drag Coefficient (Cd) vs. Mach Number (Ma) data for interpolation
MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5])
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40])

# --- II. Dynamic Physical Model Functions ---
def get_drag_coefficient(v_rel, c_sound):
    """Interpolates Cd based on Mach number."""
    mach_number = v_rel / c_sound
    Cd = np.interp(mach_number, MA_POINTS, CD_POINTS)
    return Cd

def get_air_density(z):
    """Calculates air density based on altitude (z) using standard atmosphere model."""
    h = max(0.0, z)
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0:
        return 0.0
    rho = RHO0 * (temp_ratio) ** P_TERM
    return rho

# --- III. Differential Equations (d S/d t) ---
def derivatives_of_state(S, params):
    """
    Computes the derivative of the state vector S = [x, y, z, vx, vy, vz].
    Includes Drag, Gravity, Crosswind, and Magnus Force.
    NOTE: Spin axis is along Y-axis for this simulation.
    """
    x, y, z, vx, vy, vz = S
    m, A, g, c_sound, v_wind, omega, D, cl = params

    # 1. Drag and Relative Velocity
    vx_rel = vx
    vy_rel = vy - v_wind
    vz_rel = vz
    v_rel = sqrt(vx_rel**2 + vy_rel**2 + vz_rel**2)

    if v_rel < 1e-6:
        return np.array([vx, vy, vz, 0.0, 0.0, -g])

    rho = get_air_density(z)
    Cd = get_drag_coefficient(v_rel, c_sound)

    # Drag Acceleration
    drag_coeff_factor = -(Cd * rho * A) / (2.0 * m)
    dvx_drag = drag_coeff_factor * v_rel * vx_rel
    dvy_drag = drag_coeff_factor * v_rel * vy_rel
    dvz_drag = drag_coeff_factor * v_rel * vz_rel

    # 2. Magnus Force Calculation (Spin axis along Y)
    W = np.array([0.0, omega, 0.0])         # Spin Vector: [0, omega_y, 0]
    V_rel = np.array([vx_rel, vy_rel, vz_rel]) # Relative Velocity Vector

    magnus_cross_product = np.cross(W, V_rel)
    # The cross product is: (omega*vz, 0, -omega*vx)

    # Magnus Acceleration a_M 
    if v_rel > 1e-6:
        # a_magnus = (0.5 * CL * rho * A * v_rel / M) * unit_vector(W x V_rel)
        magnus_accel_factor = (0.5 * cl * rho * A * v_rel) / M
        a_magnus = magnus_accel_factor * (magnus_cross_product / v_rel)
    else:
        a_magnus = np.array([0.0, 0.0, 0.0])

    # 3. Total Acceleration
    dvx_dt = dvx_drag + a_magnus[0] # F_Mx is non-zero
    dvy_dt = dvy_drag + a_magnus[1] # F_My is zero
    dvz_dt = dvz_drag - g + a_magnus[2] # F_Mz is non-zero (Vertical force)

    return np.array([vx, vy, vz, dvx_dt, dvy_dt, dvz_dt])

# --- IV. RK4 Integration and Trajectory Calculation (Reusable) ---
def rk4_step(S, dt, params):
    """Performs one step of the 4th-order Runge-Kutta method."""
    K1 = dt * derivatives_of_state(S, params)
    K2 = dt * derivatives_of_state(S + K1 / 2.0, params)
    K3 = dt * derivatives_of_state(S + K2 / 2.0, params)
    K4 = dt * derivatives_of_state(S + K3, params)
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

def calculate_trajectory(v0, theta_deg, v_wind, omega, dt=0.01):
    """Runs the simulation and returns the trajectory points."""
    params = (M, A, G, C_SOUND, v_wind, omega, D, MAGNUS_CL)
    theta_rad = radians(theta_deg)

    # Initial state S = [x, y, z, vx, vy, vz]
    S = np.array([0.0, 0.0, 0.0, v0 * cos(theta_rad), 0.0, v0 * sin(theta_rad)])
    trajectory = []

    # Integrate until projectile hits the ground (z <= 0)
    while S[2] >= -0.01:
        trajectory.append(S[:3].copy())
        S_new = rk4_step(S, dt, params)

        if S_new[2] < 0 and S[2] >= 0:
            break

        S = S_new

    trajectory.append(S[:3].copy())

    return np.array(trajectory)

# --- V. Visualization and Comparison Function ---
def visualize_trajectory_comparison(v0, theta, v_wind, omega_mag):
    """Calculates and plots three trajectories for comparison."""

    # 1. Baseline: No Crosswind, No Magnus (Vw=0, w=0)
    traj_baseline = calculate_trajectory(v0, theta, 0.0, 0.0)

    # 2. Crosswind Only: Vw=User Input, No Magnus (w=0)
    traj_crosswind = calculate_trajectory(v0, theta, v_wind, 0.0)

    # 3. Full Model (Y-Axis Spin): Crosswind + Magnus (Vw=User Input, w=User Input)
    traj_full = calculate_trajectory(v0, theta, v_wind, omega_mag)

    # --- Data Extraction and Plotting Setup ---
    fig = plt.figure(figsize=(14, 10))
    ax = fig.add_subplot(111, projection='3d')

    trajectories = [
        (traj_baseline, 'Baseline (No Wind, No Magnus)', 'blue', '-'),
        (traj_crosswind, f'Crosswind Only (Vw={v_wind:.1f}m/s)', 'green', '--'),
        (traj_full, f'Y-Axis Spin Model (Vw={v_wind:.1f}, ωy={omega_mag:.1f})', 'red', '-')
    ]

    impact_results = []

    for traj, label, color, style in trajectories:
        X, Y, Z = traj[:, 0], traj[:, 1], traj[:, 2]

        # Plot 3D trajectory
        ax.plot(X, Y, Z, label=label, color=color, linestyle=style, linewidth=2)
        
        # Mark impact point
        ax.scatter(X[-1], Y[-1], 0, color=color, marker='o', s=50)

        impact_results.append({
            'label': label,
            'Range': X[-1],
            'Drift': Y[-1],
            'Max_Z': np.max(Z)
        })
        
    # --- Set Labels and Title ---
    ax.set_xlabel('Range X (m) - Forward', fontsize=12)
    ax.set_ylabel('Lateral Drift Y (m) - Crosswind Direction', fontsize=12)
    ax.set_zlabel('Altitude Z (m)', fontsize=12)

    title_str = f"Trajectory Comparison with Y-Axis Spin (V₀={v0:.1f}m/s, θ={theta:.1f}°)"
    ax.set_title(title_str, fontsize=14)

    # --- Output Key Results ---
    result_info_str = "--- Simulation Results ---\n"
    for res in impact_results:
        # Note: Y-axis spin primarily affects Range (X) and Height (Z), not Drift (Y)
        result_info_str += f"{res['label']}:\n"
        result_info_str += f"  Range X: {res['Range']:.2f} m\n"
        result_info_str += f"  Max Height Z: {res['Max_Z']:.2f} m\n"
        result_info_str += f"  Final Drift Y: {res['Drift']:.2f} m\n"
        result_info_str += "\n"

    print(result_info_str)
    
    # --- Adjust View and Legend ---
    ax.view_init(elev=20, azim=-60)
    ax.legend(loc='best')
    # Display results on the plot using text2D
    ax.text2D(0.05, 0.95, result_info_str, transform=ax.transAxes, color='black', fontsize=11, verticalalignment='top')
    
    plt.show()

# --- VI. Interactive Input and Execution (Reusable) ---
def interactive_run():
    print("--- 3D Projectile Trajectory Comparison (Y-Axis Spin) ---")
    print("Enter the initial firing conditions:")
    
    try:
        v0 = float(input("Enter Initial Velocity V0 (m/s): "))
        theta = float(input("Enter Launch Angle Theta (degrees, 0-90): "))
        v_wind = float(input("Enter Crosswind Velocity V_wind (m/s, positive = +Y direction): "))
        omega_mag = float(input("Enter Spin Angular Velocity Omega (rad/s, Y-Axis): "))
        
        if v0 <= 0 or not (0 <= theta <= 90):
            print("Error: V0 must be positive and Theta must be between 0 and 90 degrees.")
            return

    except ValueError:
        print("Error: Invalid input. Please enter numerical values.")
        return

    # Run the visualization
    visualize_trajectory_comparison(v0, theta, v_wind, omega_mag)

if __name__ == "__main__":
    interactive_run()